#include "kdtree.h"
    
    
KDTree::KDTree(const std::vector<std::vector<float>>& data_points, int tree_num) 
{
    if (data_points.empty()) {
        throw std::invalid_argument("Data points cannot be empty.");
    }

    dimensions_ = data_points[0].size();
    data_points_ = data_points;

    int rows = data_points.size();
    int cols = dimensions_;

    // 创建FLANN矩阵并复制数据
    flann_data_ = std::make_unique<float[]>(rows * cols);
    for (size_t i = 0; i < rows; ++i) {
        for (size_t j = 0; j < cols; ++j) {
            flann_data_[i * cols + j] = data_points_[i][j];
        }
    }

    // Create FLANN matrix
    flann::Matrix<float> data(flann_data_.get(), rows, cols);

    // 创建KD树
    flann::KDTreeIndexParams indexParams(tree_num);
    index_ = std::make_unique<flann::Index<flann::L2<float>>>(data, indexParams);
    index_->buildIndex();
}

// 搜索最近点的方法
std::vector<float> KDTree::findNearestPoint(const std::vector<float>& query_point, int& index) const 
{
    if (query_point.size() != dimensions_) {
        std::cout<<query_point.size()<<" != "<<dimensions_<<std::endl;
        throw std::invalid_argument("Query point dimensions must match data point dimensions.");
    }

    flann::Matrix<float> query(const_cast<float*>(query_point.data()), 1, dimensions_);

    // 结果存储
    std::vector<std::vector<int>> indices;
    std::vector<std::vector<float>> dists;

    // 执行knn搜索
    flann::SearchParams param;
    param.checks = -1;
    param.eps = 0;
    param.sorted = true;
    param.cores = 0;
    index_->knnSearch(query, indices, dists, 1, param);

    // for(int i=0; i<query.rows; i++)
    // {
    //     for(int j=0; j<query.cols; j++)
    //     {
    //         std::cout<<query[i][j]<<", ";
    //     }
    //     std::cout<<std::endl;
    // }
    // std::cout<<"idx:"<<indices[0][0]<<std::endl;

    // 通过索引获取最近点
    index = indices[0][0];
    return data_points_[indices[0][0]];
}

